"""金融知识图谱集成模块"""
import networkx as nx
import json
from typing import Dict, List, Any, Optional
from fin_senti_entity_platform.utils.logger import get_logger
from fin_senti_entity_platform.utils.constants import FINANCIAL_ENTITY_TYPES

logger = get_logger(__name__)

class FinancialKnowledgeGraph:
    """金融知识图谱类，用于存储和查询金融实体间关系"""
    
    def __init__(self):
        self.graph = nx.DiGraph()
        self._load_knowledge_base()
    
    def _load_knowledge_base(self):
        """加载金融领域知识库"""
        try:
            # 加载预定义的金融实体关系
            with open('data/knowledge/financial_relations.json', 'r', encoding='utf-8') as f:
                relations = json.load(f)
                
            # 构建知识图谱
            for relation in relations:
                self.add_relation(
                    relation['source'],
                    relation['target'],
                    relation['relation_type'],
                    relation.get('attributes', {})
                )
                
            logger.info(f"成功加载金融知识图谱，包含 {len(self.graph.nodes)} 个实体和 {len(self.graph.edges)} 条关系")
        except Exception as e:
            logger.error(f"加载金融知识图谱失败: {str(e)}")
    
    def add_relation(self, source: str, target: str, relation_type: str, attributes: Optional[Dict] = None):
        """添加实体间关系"""
        if attributes is None:
            attributes = {}
        
        self.graph.add_node(source, entity_type=attributes.get('source_type'))
        self.graph.add_node(target, entity_type=attributes.get('target_type'))
        self.graph.add_edge(source, target, relation_type=relation_type, **attributes)
    
    def get_related_entities(self, entity: str, relation_type: Optional[str] = None) -> List[Dict[str, Any]]:
        """获取与指定实体相关的所有实体"""
        result = []
        
        if entity not in self.graph:
            return result
        
        for neighbor in self.graph.neighbors(entity):
            edge_data = self.graph.get_edge_data(entity, neighbor)
            if relation_type is None or edge_data['relation_type'] == relation_type:
                result.append({
                    'entity': neighbor,
                    'relation_type': edge_data['relation_type'],
                    'attributes': {k: v for k, v in edge_data.items() if k != 'relation_type'}
                })
        
        return result